import torch
from torch import nn


class PositionalEncoder(torch.nn.Module):
    def __init__(self, num_hiddens, dropout, max_len=1000):
        """
        在位置嵌入矩阵  P  中，行代表词元在序列中的位置，列代表位置编码的不同维度。
        Args:
            num_hiddens: 位置编码的维度
            dropout:
            max_len: 要编码序列的长度
        """
        super(PositionalEncoder, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的 `P`
        self.P = torch.zeros((1, max_len, num_hiddens))

        #  X : (max_len,16)
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)

        # 切片x：：y：从x开始每隔y-1个取一个
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)
